import dask.dataframe as dd
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional
import json

class FactorStorage:
    """因子存储管理（支持Parquet格式分块存储）"""
    
    def __init__(self, storage_path: str = "data/factors"):
        self.base_path = Path(storage_path)
        self.meta_path = self.base_path / "_metadata"
        self.base_path.mkdir(parents=True, exist_ok=True)
        
    def save(self, factor_name: str, data: dd.DataFrame, 
            time_range: Dict[str, datetime], 
            overwrite: bool = False) -> None:
        """保存因子数据"""
        path = self._get_file_path(factor_name, time_range)
        if path.exists() and not overwrite:
            raise FileExistsError(f"因子数据已存在: {path}")
            
        data.to_parquet(
            path,
            engine="pyarrow",
            partition_on=["date"],
            write_metadata_file=True
        )
        self._update_metadata(factor_name, time_range)
        
    def load(self, factor_name: str, 
            start: Optional[datetime] = None,
            end: Optional[datetime] = None) -> dd.DataFrame:
        """加载因子数据"""
        meta = self._read_metadata(factor_name)
        if not meta:
            raise FileNotFoundError(f"未找到因子元数据: {factor_name}")
            
        filters = []
        if start:
            filters.append(('date', '>=', start.date().isoformat()))
        if end:
            filters.append(('date', '<=', end.date().isoformat()))
            
        return dd.read_parquet(
            self._get_file_path(factor_name, meta["time_range"]),
            filters=filters if filters else None
        )
        
    def _get_file_path(self, factor_name: str, time_range: Dict) -> Path:
        """生成存储路径"""
        start_str = time_range["start"].date().isoformat()
        end_str = time_range["end"].date().isoformat()
        return self.base_path / f"{factor_name}_{start_str}_{end_str}"
        
    def _update_metadata(self, factor_name: str, time_range: Dict) -> None:
        """更新元数据文件"""
        meta = self._read_all_metadata()
        meta[factor_name] = {
            "time_range": {
                "start": time_range["start"].isoformat(),
                "end": time_range["end"].isoformat()
            },
            "storage_path": str(self._get_file_path(factor_name, time_range))
        }
        with open(self.meta_path, 'w') as f:
            json.dump(meta, f)
            
    def _read_metadata(self, factor_name: str) -> Optional[Dict]:
        """读取单个因子元数据"""
        meta = self._read_all_metadata()
        return meta.get(factor_name)
        
    def _read_all_metadata(self) -> Dict:
        """读取全部元数据"""
        if not self.meta_path.exists():
            return {}
        with open(self.meta_path) as f:
            return json.load(f)